package main

import (
	"errors"
	"fmt"
	"github.com/fatih/color"
	"net/http"
	"os"
	"sync"
)

// ParallelGetTask 多线程下载任务类，存放一个多线程下载任务的参数和状态量
type ParallelGetTask struct {
	// 文件的下载链接
	Url string
	// 文件的最终保存位置
	FilePath string
	// 下载并发数
	Concurrent int
	// 下载文件的总大小
	TotalSize int64
	// 全部的下载分片任务参数列表
	ShardTaskList []*ShardTask
}

// NewParallelGetTask 构造函数
func NewParallelGetTask(url, filePath string, concurrent int) *ParallelGetTask {
	return &ParallelGetTask{
		// 参数赋值
		Url:        url,
		FilePath:   filePath,
		Concurrent: concurrent,
		// 初始化状态量
		TotalSize:     0,
		ShardTaskList: make([]*ShardTask, 0),
	}
}

// 发送HEAD请求获取待下载文件的大小
func (task *ParallelGetTask) getLength() error {
	// 发送HEAD请求
	response, e := http.Head(task.Url)
	if e != nil {
		color.Red("发送HEAD请求出错！")
		return e
	}
	// 判断状态码是否正确，不正确说明Head不被允许，切换为Get重试
	if response.StatusCode >= 300 {
		color.HiYellow("不支持HEAD请求，状态码：%d，使用GET重试...", response.StatusCode)
		response, e = http.Get(task.Url)
		if e != nil {
			color.Red("发送GET请求获取大小出错！")
			return e
		}
		// 最终直接关闭响应体，不进行读取
		defer func() {
			_ = response.Body.Close()
		}()
		// 再次检查状态码，若不正确则返回错误
		if response.StatusCode >= 300 {
			color.Red("发送GET请求获取大小出错！状态码：%d", response.StatusCode)
			return errors.New(fmt.Sprintf("状态码不正确：%d", response.StatusCode))
		}
	}
	// 检查是否支持部分请求
	if response.Header.Get("Accept-Ranges") != "bytes" {
		color.Red("该请求不支持部分获取，无法分片下载！")
		return errors.New("不支持分片获取")
	}
	// 读取并设定长度
	task.TotalSize = response.ContentLength
	if task.TotalSize <= 0 {
		color.Red("无法获取内容长度！")
		return errors.New("不能获取内容长度")
	}
	return nil
}

// 根据待下载文件的大小和设定的并发数，创建每个分片任务对象
func (task *ParallelGetTask) allocateTask() {
	// 如果并发数大于总大小，则进行调整
	if int64(task.Concurrent) > task.TotalSize {
		task.Concurrent = int(task.TotalSize)
	}
	// 开始计算每个分片的下载范围
	eachSize := task.TotalSize / int64(task.Concurrent)
	// 创建任务对象
	for i := 0; i < task.Concurrent; i++ {
		task.ShardTaskList = append(task.ShardTaskList, NewShardTask(task.Url, i+1, task.FilePath, int64(i)*eachSize, int64(i+1)*eachSize-1))
	}
	// 处理末尾部分
	if task.TotalSize%int64(task.Concurrent) != 0 {
		task.ShardTaskList[task.Concurrent-1].RangeEnd = task.TotalSize - 1
	}
}

// 创建一个与目标下载文件大小一样的空白的文件
func (task *ParallelGetTask) createFile() error {
	file, e := os.OpenFile(task.FilePath, os.O_WRONLY|os.O_CREATE, 0755)
	if e != nil {
		color.Red("创建文件出错！")
		return e
	}
	defer func() {
		_ = file.Close()
	}()
	// Truncate能够调整文件到指定大小
	// 若文件大小小于给定大小，则会使用空白字符填充并扩充至给定大小
	// 否则截断文件至给定大小，丢弃多余部分
	e = file.Truncate(task.TotalSize)
	if e != nil {
		color.Red("调整文件大小出错！")
		return e
	}
	return nil
}

// 根据任务列表进行多线程分片下载操作
func (task *ParallelGetTask) downloadShard() {
	// 创建线程组
	waitGroup := &sync.WaitGroup{}
	// 开始执行全部分片下载线程
	for _, task := range task.ShardTaskList {
		go task.DoShardGet(waitGroup)
		waitGroup.Add(1)
	}
	// 等待全部下载完成
	waitGroup.Wait()
}

// Run 开始执行整个分片多线程下载任务
func (task *ParallelGetTask) Run() error {
	// 获取文件大小
	e := task.getLength()
	if e != nil {
		color.Red("%s", e)
		return e
	}
	color.HiYellow("已获取到下载文件大小：%d字节", task.TotalSize)
	// 分配任务
	task.allocateTask()
	color.HiYellow("已完成分片任务分配，共计%d个任务", len(task.ShardTaskList))
	// 创建空白文件
	e = task.createFile()
	if e != nil {
		color.Red("%s", e)
		return e
	}
	color.HiYellow("已预创建目标下载文件！")
	// 开启进度输出
	printProcess(task)
	// 开始下载文件
	task.downloadShard()
	color.Green("\n分片下载任务完成！")
	return nil
}